import streamlit as st
import pywencai
import requests
import json
import pandas as pd
import akshare as ak
from datetime import datetime

# ==================== 配置常量 ====================
CONCEPT_INDEX_QUERY = "同花顺概念指数"
CACHE_TTL = 3600  # 1小时缓存
REQUEST_TIMEOUT = 10
COLUMN_CONFIG = {
    'code': {'display': '代码', 'width': 3, 'sortable': True},
    '指数简称': {'display': '简称', 'width': 4, 'sortable': True},
    '涨跌幅': {'display': '涨幅', 'width': 3, 'sortable': True, 'format': '{:.2%}', 'default_sort': True}
}

# ==================== 数据获取函数 ====================
@st.cache_data(ttl=CACHE_TTL)
def get_concept_index() -> pd.DataFrame:
    """获取概念指数数据并缓存"""
    try:
        return pywencai.get(
            query=CONCEPT_INDEX_QUERY,
            query_type="zhishu",
            sort_order='desc',
            loop=True
        )
    except Exception as e:
        st.error(f"获取概念指数数据失败: {str(e)}")
        return pd.DataFrame()

@st.cache_data(ttl=CACHE_TTL)
def get_trade_date() -> str:
    """获取最近交易日"""
    try:
        df = ak.tool_trade_date_hist_sina()
        df['trade_date'] = pd.to_datetime(df['trade_date'], errors='coerce')
        df = df.dropna(subset=['trade_date'])
        df = df[df['trade_date'] <= pd.to_datetime(datetime.now().date())]
        
        if df.empty:
            raise ValueError("没有有效的交易日数据")
            
        return df.sort_values('trade_date', ascending=False)['trade_date'].iloc[0].strftime("%Y%m%d")
    except Exception as e:
        st.warning(f"获取交易日失败: {str(e)}，使用当日日期")
        return datetime.now().strftime("%Y%m%d")

# ==================== 组件函数 ====================
def init_session_state():
    """初始化session状态"""
    if 'sort_column' not in st.session_state:
        st.session_state.sort_column = next(
            (col for col, config in COLUMN_CONFIG.items() if config.get('default_sort', False)),
            '涨跌幅'
        )
    if 'sort_ascending' not in st.session_state:
        st.session_state.sort_ascending = False
    if 'selected_code' not in st.session_state:
        st.session_state.selected_code = None

def render_sortable_header(df: pd.DataFrame) -> pd.DataFrame:
    """渲染可排序表头并返回排序后的数据"""
    header_cols = st.columns([config['width'] for config in COLUMN_CONFIG.values()])
    
    for i, (col_key, config) in enumerate(COLUMN_CONFIG.items()):
        if not config.get('sortable', False):
            continue
            
        with header_cols[i]:
            is_active = st.session_state.sort_column == col_key
            arrow = "▼" if is_active and not st.session_state.sort_ascending else "▲"
            btn_label = f"{config['display']} {arrow}" if is_active else config['display']
            
            if st.button(btn_label, key=f"sort_{col_key}"):
                if is_active:
                    st.session_state.sort_ascending = not st.session_state.sort_ascending
                else:
                    st.session_state.sort_column = col_key
                    st.session_state.sort_ascending = False
                st.rerun()
    
    return df.sort_values(
        st.session_state.sort_column,
        ascending=st.session_state.sort_ascending
    )


def format_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """格式化DataFrame显示"""
    df = df.copy()
    for col, config in COLUMN_CONFIG.items():
        if col in df.columns and 'format' in config:
            try:
                if pd.api.types.is_numeric_dtype(df[col]):
                    #df[col] = df[col].apply(lambda x: config['format'].format(x))
                    df[col] = df[col].apply(lambda x: f"{x:,.2f}%")
            except Exception:
                pass
    return df

def render_concept_row(row: pd.Series):
    """渲染单个概念指数行"""
    cols = st.columns([config['width'] for config in COLUMN_CONFIG.values()])
    
    with cols[0]:  # 代码列
        if st.button(
            row['code'],
            key=f"btn_{row['code']}",
            help="点击查看成分股",
            use_container_width=True
        ):
            st.session_state.selected_code = row['code']
            st.rerun()
    
    with cols[1]:  # 名称列
        st.markdown(f'<div style="font-size:14px">{row["指数简称"]}</div>', unsafe_allow_html=True)
    
    with cols[2]:  # 涨跌幅列
        change = float(row['涨跌幅'].rstrip('%')) if isinstance(row['涨跌幅'], str) else row['涨跌幅']
        color = "#ff4b4b" if change > 0 else "#0db14b" if change < 0 else "gray"
        display_value = f"{change:.2f}%" if isinstance(change, float) else row['涨跌幅']
        st.markdown(
            f'<div style="color:{color};font-size:14px;text-align:right">{display_value}</div>',
            unsafe_allow_html=True
        )

def safe_float(value, default=0.0):
    """安全转换为浮点数"""
    try:
        return float(str(value).strip()) if str(value).strip() else default
    except ValueError:
        return default
    
def show_stock_list(code: str):
    """显示成分股列表"""
    url = f"https://d.10jqka.com.cn/v2/blockrank/{code}/199112/d1000.js"
    headers = {
        'Referer': 'http://q.10jqka.com.cn/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
                      '(KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36'
    }

    try:
        with st.spinner("正在加载成分股..."):
            response = requests.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
            response.raise_for_status()
            
            json_str = response.text.split('(', 1)[1].rsplit(')', 1)[0]
            data = json.loads(json_str)
            stock_list = data.get('items', [])
            
            if not stock_list:
                st.warning("未找到相关个股数据")
                return
                
            # 创建DataFrame
            stocks_df = pd.DataFrame([
                (
                    s.get('5', '').zfill(6),
                    s.get('55', ''),
                    safe_float(s.get('8', 0)),
                    safe_float(s.get('199112', 0))
                ) for s in stock_list],
                columns=['股票代码', '股票名称', '最新价', '涨跌幅']
            )
            
            # 定义颜色样式函数
            def color_change(val):
                color = 'red' if val > 0 else 'green' if val < 0 else 'gray'
                return f'color: {color}'
            
            # 应用样式
            styled_df = stocks_df.style.applymap(color_change, subset=['涨跌幅'])\
                                      .format({
                                          '最新价': '{:.2f}',
                                          '涨跌幅': '{:.2f}%'
                                      })
            
            # 显示带样式的数据框
            st.dataframe(
                styled_df,
                use_container_width=True,
                hide_index=True,
                column_config={
                    "股票代码": st.column_config.TextColumn(width="small"),
                    "股票名称": st.column_config.TextColumn(width="medium"),
                    "最新价": st.column_config.NumberColumn(format="%.2f"),
                    "涨跌幅": st.column_config.NumberColumn(format="%.2f%%")
                }
            )
            
    except Exception as e:
        st.error(f"获取成分股数据失败: {str(e)}")


def concept_index_section(cur_date_str: str):
    """概念指数列表部分"""
    df = get_concept_index()
    if df.empty:
        st.warning("未能获取概念指数数据")
        return
    
    # 列名重映射和筛选
    df = df.rename(columns={
        '指数简称': '指数简称',
        'code': 'code',
        f'指数@涨跌幅:前复权[{cur_date_str}]': '涨跌幅'
    })[list(COLUMN_CONFIG.keys())]
    
    # 排序和格式化
    df = render_sortable_header(df)
    df = format_dataframe(df)
    
    # 渲染表格
    for _, row in df.iterrows():
        render_concept_row(row)
    
    st.caption(f"总共显示 {len(df)} 个概念指数")

def app():
    """主应用函数"""
    # st.set_page_config(
    #     page_title="同花顺概念指数分析",
    #     layout="wide",
    #     page_icon="📊"
    # )
    st.title("📊 同花顺概念指数分析")
    
    # 初始化状态和样式
    init_session_state()
    st.markdown("""
    <style>
    .stButton>button {
        min-height: auto;
        padding-top: 2px;
        padding-bottom: 2px;
    }
    .stock-code {
        font-family: monospace;
    }
    .positive-change {
        color: #ff4b4b;
    }
    .negative-change {
        color: #0db14b;
    }
    .neutral-change {
        color: gray;
    }
    .dataframe th {
        font-size: 12px !important;
    }
    .dataframe td {
        font-size: 12px !important;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # 获取交易日
    cur_date_str = get_trade_date()
    
    # 创建两列布局
    col1, col2 = st.columns([3, 7])

    # 左侧概念指数列表
    with col1:
        with st.expander("📊 概念指数列表", expanded=True):
            concept_index_section(cur_date_str)

    # 右侧成分股展示
    with col2:
        if st.session_state.selected_code:
            with st.expander(f"📈 成分股列表 - {st.session_state.selected_code}", expanded=True):
                show_stock_list(st.session_state.selected_code)
        else:
            with st.expander("ℹ️ 成分股列表", expanded=True):
                st.info("请从左侧选择概念指数查看其成分股")
                st.markdown("""
                <div style="text-align: center; padding: 20px; border: 1px dashed #ccc; border-radius: 5px;">
                    <p>请从左侧选择概念指数</p>
                    <p>👈 点击任意指数代码查看其成分股</p>
                </div>
                """, unsafe_allow_html=True)

    # 添加页脚
    st.markdown("---")
    st.caption(f"数据来源: 同花顺 | 更新时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

if __name__ == "__main__":
    app()
